import shodan

from export.export import export_to_excel
from .base import IScan
from config.settings import SHODAN_API, MIN_NETWORK_NUM, SHODAN_SAVE_PATH
import json
import os


class ShodanScan(IScan):
    name = 'Shodan'

    def __init__(self):
        self.register(SHODAN_API)
        self.shodan_info = []
        self.shodan_network = []
        self.export_results = []
        self.ip = ''

    @property
    def fields(self):
        # vulns（漏洞）信息单独处理
        return [
            'port',
            'timestamp',
            'hostnames',
            'org',
            'isp',
            'location.country_name',
            'location.country_code',
            'location.city',
            'os',
            'asn',
            'data',
            'transport',
            'domains',
            'http.server',
            'http.title'
        ]

    @property
    def columns_dict(self):
        return {
            'ip': 'IP',
            'port': '端口',
            'timestamp': '数据时间',
            'hostnames': '主机名',
            'org': '组织名',
            'isp': 'ISP',
            'location.country_name': '国家',
            'location.country_code': 'CODE',
            'location.city': '城市',
            'os': '操作系统',
            'asn': 'ASN区域',
            'transport': '协议',
            'data': 'banner信息',
            'domains': '绑定域名',
            'http.server': 'web服务器',
            'http.title': 'web页面标题'
        }
    @property
    def inputs(self):
        return {
            'ip': '',
        }

    @property
    def outputs(self):
        return {
            'export_results': self.export_results
        }

    def register(self, api):
        self.shodan = shodan.Shodan(api)

    def get_ip_info(self):
        for ip in self.ip:
            self.shodan_info.append(
                self._ip_info(ip)
            )

    def save(self, ip):
        if '/' in ip:
            ip = ip.replace('/', '_')
        with open(f'{SHODAN_SAVE_PATH}/{ip}.json', 'w+') as f:
            json.dump(self.result, f)

    def cache(self, ip: str):
        if '/' in ip:
            ip = ip.replace('/', '_')
        if os.path.exists(f'{SHODAN_SAVE_PATH}/{ip}.json'):
            with open(f'{SHODAN_SAVE_PATH}/{ip}.json', 'r') as f:
                self.result = json.load(f)
            if self.result:
                return True
            else:
                self.result = {}
                return False

    def _ip_info(self, ip):
        self.ip = ip
        try:
            if self.cache(ip):
                return self.result
            info = self.shodan.host(ip)
            self.result = info
            self.save(ip)
        except Exception as e:
            print(e)
            info = {}
            self.result = info
        return info

    def get_network_info(self):
        if isinstance(self.network, list):
            for network in self.network:
                self.shodan_network.append(
                    self._network_info(network)
                )
        elif isinstance(self.network, dict):
            self.network = [key for key in self.network if len(self.network[key]) > MIN_NETWORK_NUM]
            self.get_network_info()
        else:
            self.shodan_network = self._network_info(self.network)

    def _network_info(self, network):
        try:
            if self.cache(network):
                return self.result
            info = self.shodan.search(f'net:{network}')
            self.result = info
            self.save(network)
        except:
            info = {}
        return info

    def parse(self):
        export_results = []

        obj = self.result
        if not obj:
            return
        if 'data' in obj:
            for result in obj['data']:
                export_result = {}
                for field in self.fields:
                    value = self.get_fields(result, field)
                    export_result[self.columns_dict[field]] = value
                if 'vulns' in result:
                    cves = result['vulns'].keys()
                    cves = ','.join(cves)
                else:
                    cves = ''
                export_result['漏洞信息'] = cves
                export_result['engine'] = 'shodan'
                export_result['source'] = f'{SHODAN_SAVE_PATH}/{self.ip}.json'
                export_results.append(export_result)
        self.export_results = export_results
        self.outputs['export_results'] = export_results

    def export(self):
        return export_to_excel(self.export_results, sheet_name='IP信息(shodan)', xlsx_name=f'{self.ip}')

    def run(self):
        self.get_ip_info()
        self.get_network_info()


class Shodan:
    def __init__(self, ip_list: list, query='shodan_search'):
        self.ip_list = ip_list
        self.scan = ShodanScan()
        self.results = []
        self.query = query

    def run(self, save_path):
        for ip in self.ip_list:
            print(f'- {ip} 正在使用shodan查询...')
            self.scan._ip_info(ip)
            self.scan.parse()
            self.results.extend(self.scan.outputs['export_results'])

        self.export(save_path)

    def export(self, save_path=None):
        return export_to_excel(self.results, sheet_name='IP信息(shodan)', xlsx_name=f'{self.query}', save_path=save_path)